# 测试hanlp的命名实体识别方法，使用感知机
from pyhanlp import *
import os
from util import ensure_data

# 清华大学语料库
PKU98 = ensure_data("pku98", "http://file.hankcs.com/corpus/pku98.zip")
PKU199801 = os.path.join(PKU98, "199801.txt")
PKU199801_TRAIN = os.path.join(PKU98, "199801-train.txt")
PKU199801_TEST = os.path.join(PKU98, "199801-test.txt")
POS_MODEL = os.path.join(PKU98, "pos.bin")
NER_MODEL = os.path.join(PKU98, "ner.bin")

# 分词器
PerceptronSegmenter = JClass("com.hankcs.hanlp.model.perceptron.PerceptronSegmenter")
# 词性标注器
PerceptronPOSTagger = JClass("com.hankcs.hanlp.model.perceptron.PerceptronPOSTagger")
NERTrainer = JClass("com.hankcs.hanlp.model.perceptron.NERTrainer")
PerceptronNERecognizer = JClass(
    "com.hankcs.hanlp.model.perceptron.PerceptronNERecognizer"
)

# model_path = ""
model_path = "demo/testdata/pku98/ner.bin"

# 构建命名实体识别器
trainer = NERTrainer()
if model_path == "":  # 没有传入模型路径就重新下载数据集并训练
    recognizer = PerceptronNERecognizer(
        # 通过下载的语料训练模型，保存到ner.bin
        trainer.train(PKU199801_TRAIN, NER_MODEL).getModel()
    )
else:
    recognizer = PerceptronNERecognizer(model_path)

# 将上述三个处理器合并
analyzer = PerceptronLexicalAnalyzer(
    PerceptronSegmenter(), PerceptronPOSTagger(), recognizer
)

# 识别并输出
print(analyzer.analyze("华北电力公司董事长谭旭光和秘书胡花蕊来到美国纽约现代艺术博物馆参观。"))
